from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import pycocotools.coco as coco
from pycocotools.cocoeval import COCOeval
import numpy as np
import json
import os

import torch.utils.data as data


class COCOHP(data.Dataset):
    """
    The order of joints:
        KEYPOINT_DICT = {
        'nose': 0,
        'left_eye': 1,
        'right_eye': 2,
        'left_ear': 3,
        'right_ear': 4,
        'left_shoulder': 5,
        'right_shoulder': 6,
        'left_elbow': 7,
        'right_elbow': 8,
        'left_wrist': 9,
        'right_wrist': 10,
        'left_hip': 11,
        'right_hip': 12,
        'left_knee': 13,
        'right_knee': 14,
        'left_ankle': 15,
        'right_ankle': 16
        }
    """

    num_classes = 1
    num_joints = 17
    default_resolution = [512, 512]
    mean = np.array([0.40789654, 0.44719302, 0.47026115], dtype=np.float32).reshape(
        1, 1, 3
    )
    std = np.array([0.28863828, 0.27408164, 0.27809835], dtype=np.float32).reshape(
        1, 1, 3
    )
    flip_idx = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16]]

    def __init__(self, opt, split, sp=False):
        super(COCOHP, self).__init__()
        self.edges = [
            [0, 1],
            [0, 2],
            [1, 3],
            [2, 4],
            [4, 6],
            [3, 5],
            [5, 6],
            [5, 7],
            [7, 9],
            [6, 8],
            [8, 10],
            [6, 12],
            [5, 11],
            [11, 12],
            [12, 14],
            [14, 16],
            [11, 13],
            [13, 15],
        ]

        self.acc_idxs = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
        self.data_dir = os.path.join(opt.data_dir, "coco")
        self.img_dir = os.path.join(self.data_dir, "{}2017".format(split))
        if split == "test":
            self.annot_path = os.path.join(
                self.data_dir, "annotations", "image_info_test-dev2017.json"
            ).format(split)
        else:
            self.annot_path = os.path.join(
                self.data_dir, "annotations", "person_keypoints_{}2017.json"
            ).format(split)
        self.max_objs = 32
        self._data_rng = np.random.RandomState(123)
        self._eig_val = np.array([0.2141788, 0.01817699, 0.00341571], dtype=np.float32)
        self._eig_vec = np.array(
            [
                [-0.58752847, -0.69563484, 0.41340352],
                [-0.5832747, 0.00994535, -0.81221408],
                [-0.56089297, 0.71832671, 0.41158938],
            ],
            dtype=np.float32,
        )
        self.split = split
        self.opt = opt

        print("==> initializing coco 2017 {} data.".format(split))
        self.coco = coco.COCO(self.annot_path)
        image_ids = self.coco.getImgIds()

        if split == "train":
            self.images = []
            for img_id in image_ids:
                idxs = self.coco.getAnnIds(imgIds=[img_id])
                if len(idxs) > 0:
                    self.images.append(img_id)
        else:
            self.images = image_ids
        self.num_samples = len(self.images)
        print("Loaded {} {} samples".format(split, self.num_samples))

    def _to_float(self, x):
        return float("{:.2f}".format(x))

    def convert_eval_format(self, all_bboxes):
        # import pdb; pdb.set_trace()
        detections = []
        for image_id in all_bboxes:
            for cls_ind in all_bboxes[image_id]:
                category_id = 1
                for dets in all_bboxes[image_id][cls_ind]:
                    bbox = dets[:4]
                    bbox[2] -= bbox[0]
                    bbox[3] -= bbox[1]
                    score = dets[4]
                    bbox_out = list(map(self._to_float, bbox))
                    keypoints = (
                        np.concatenate(
                            [
                                np.array(dets[5:39], dtype=np.float32).reshape(-1, 2),
                                np.ones((17, 1), dtype=np.float32),
                            ],
                            axis=1,
                        )
                        .reshape(51)
                        .tolist()
                    )
                    keypoints = list(map(self._to_float, keypoints))

                    detection = {
                        "image_id": int(image_id),
                        "category_id": int(category_id),
                        "bbox": bbox_out,
                        "score": float("{:.2f}".format(score)),
                        "keypoints": keypoints,
                    }
                    detections.append(detection)
        return detections

    def __len__(self):
        return self.num_samples

    def save_results(self, results, save_dir):
        json.dump(
            self.convert_eval_format(results),
            open("{}/results.json".format(save_dir), "w"),
        )

    def run_eval(self, results, save_dir):
        self.save_results(results, save_dir)
        coco_dets = self.coco.loadRes("{}/results.json".format(save_dir))
        coco_eval = COCOeval(self.coco, coco_dets, "keypoints")
        coco_eval.evaluate()
        coco_eval.accumulate()
        coco_eval.summarize()
        coco_eval = COCOeval(self.coco, coco_dets, "bbox")
        coco_eval.evaluate()
        coco_eval.accumulate()
        coco_eval.summarize()
